/* 
 *  Incub - incubate your baby robot
 *  Copyright (C) 2008 Stanislaw Szymczyk
 *  e-mail: sszymczy@gmail.com   www: http://cogville.cog-lab.com
 *
 *  This file is part of Incub.
 *
 *  Incub is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  Incub is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with Incub.  If not, see <http://www.gnu.org/licenses/>.
 * 
 */
 
#include <algorithm>

#include "Physics.h"
#include "PhysicalJoint.h"
#include "PhysicalSensor.h"
#include "PhysicalShape.h"
#include "Plane.h"

using std::map;
using std::pair;
using std::vector;
using std::set;

World* Physics::root;
dWorldID Physics::dynWorld;
dSpaceID Physics::colSpace;
dJointGroupID Physics::contactGroup;
map<Entity*, dBodyID> Physics::bodies;
map<dBodyID, Vector3D> Physics::shifts;
map<Joint*, dJointID> Physics::joints;
set<pair<dBodyID, dBodyID> > Physics::jointBodies;
map<dJointFeedback*, pair<Shape*, Shape*> > Physics::feedbacks;
map<Shape*, Vector3D> Physics::forces;

int Physics::counter = 0;

void rpyToQ(dQuaternion rotQ, float roll, float pitch, float yaw) {
	// calculate rotation quaternion from roll pitch yaw (XYZ convention)
	double factor = (2*M_PI)/(2*360);
	double sr = sin( factor*roll ); double cr = cos( factor*roll );
	double sp = sin( factor*pitch ); double cp = cos( factor*pitch );
	double sy = sin( factor*yaw ); double cy = cos( factor*yaw );
    rotQ[0] = cr * cp * cy + sr * sp * sy;
    rotQ[1] = sr * cp * cy - cr * sp * sy;
    rotQ[2] = cr * sp * cy + sr * cp * sy;
    rotQ[3] = cr * cp * sy - sr * sp * cy;
}

void qtoRpy(double *rpy, dQuaternion q) {
	double factor = 360.0/(2*M_PI);
	rpy[0] = factor*atan2(2*q[2]*q[3] + 2*q[0]*q[1], q[3]*q[3] - q[2]*q[2] - q[1]*q[1] + q[0]*q[0]);;
	rpy[1] = -factor*asin(2*q[1]*q[3] - 2*q[0]*q[2]);
	rpy[2] = factor*atan2(2*q[1]*q[2] + 2*q[0]*q[3], q[1]*q[1] + q[0]*q[0] - q[3]*q[3] - q[2]*q[2]);
}

void getLocalPose(dVector3 localPos, dMatrix3 localRot, Transformable *trans) {
	dQuaternion rotQ;
	rpyToQ(rotQ, trans->getRoll(), trans->getPitch(), trans->getYaw());
	dQtoR(rotQ, localRot);

	localPos[0] = trans->getX();
	localPos[1] = trans->getY();
	localPos[2] = trans->getZ();
}

void getGlobalPose(dVector3 newGlobalPos, dMatrix3 newGlobalRot, dVector3 globalPos, dMatrix3 globalRot, dVector3 localPos, dMatrix3 localRot) {
    dVector3 localPosRotated;
    dMULTIPLY0_331 (localPosRotated,globalRot,localPos);
    
    newGlobalPos[0] = globalPos[0] + localPosRotated[0];
    newGlobalPos[1] = globalPos[1] + localPosRotated[1];
    newGlobalPos[2] = globalPos[2] + localPosRotated[2];

    dMULTIPLY0_333 (newGlobalRot,globalRot,localRot);
}

dGeomID Physics::initShape(Entity* entity, dBodyID body, dSpaceID colSpace, dMass* mass, Shape* shape, dVector3 globalPos, dMatrix3 globalRot) {
	dGeomID geom;
	dMass shapeMass;

	dMatrix3 localRot;
	dVector3 localPos;
	getLocalPose(localPos, localRot, shape);
  
    dVector3 newGlobalPos;
    dMatrix3 newGlobalRot;
	getGlobalPose(newGlobalPos, newGlobalRot, globalPos, globalRot, localPos, localRot);
	
	if(!dynamic_cast<Plane*>(shape) && !entity->getStill()) colSpace = 0;

	Material* mat = shape->getMaterial();
		
	geom = dynamic_cast<GenericPhysicalShape*>(shape)->createGeom(colSpace);
	dynamic_cast<GenericPhysicalShape*>(shape)->setMass(&shapeMass, mat->getDensity());
	
	if(dynamic_cast<Plane*>(shape)) {
		Plane* plane = dynamic_cast<Plane*>(shape);
		// since plane in non-placeable geom in ODE, we have to
		// calculate plane equation parameters from global position and rotation
		dVector3 normal;
		dSetZero(normal, 3);
		normal[1] = 1;
	    dVector3 normalRotated;
	    dMULTIPLY0_331 (normalRotated,newGlobalRot,normal);
		dReal d = dDOT(newGlobalPos, normalRotated);
		
		dGeomPlaneSetParams(geom, normalRotated[0], normalRotated[1], normalRotated[2], d);
	}

	dMassTranslate(&shapeMass, localPos[0], localPos[1], localPos[2]);	
	dMassRotate(&shapeMass, localRot);
	dMassAdd(mass, &shapeMass);
//	cout << "shapeMass center: " << shapeMass->c[0] << ", " << shapeMass->c[1] << ", " << -shapeMass->c[2] << endl;
    
	if(!dynamic_cast<Plane*>(shape)) {
		if(!entity->getStill()) {
			// not still - assigning local transforms
			dGeomID transform = dCreateGeomTransform(Physics::colSpace);
			dGeomTransformSetInfo(transform, 1);
			dGeomSetPosition(geom, localPos[0], localPos[1], localPos[2]);
			dGeomSetRotation(geom, localRot);
			dGeomTransformSetGeom (transform,geom);
			dGeomSetBody(transform, body);
		} else {
			// still entity - assigning global transforms
			dGeomSetPosition(geom, newGlobalPos[0], newGlobalPos[1], newGlobalPos[2]);
			dGeomSetRotation(geom, newGlobalRot);
			dGeomSetBody(geom, 0);	
		}
	} else {
		dGeomSetBody(geom, 0);	
	}

	dGeomSetData(geom, shape);

	return geom;
}

void Physics::initJoint(Joint* joint, dBodyID body1, dBodyID body2, dVector3 globalPos, dMatrix3 globalRot) {
	GenericPhysicalJoint *physicalJoint = dynamic_cast<GenericPhysicalJoint*>(joint);
	dJointID jointID = physicalJoint->init(dynWorld, body1, body2, globalPos, globalRot);
	joints[joint] = jointID;
	jointBodies.insert(std::make_pair(body1, body2));
}

dBodyID Physics::initEntity(Entity* entity, dVector3 globalPos, dMatrix3 globalRot) {
	dMatrix3 localRot;
	dVector3 localPos;
	getLocalPose(localPos, localRot, entity);
   
    dVector3 newGlobalPos;
    dMatrix3 newGlobalRot;
	getGlobalPose(newGlobalPos, newGlobalRot, globalPos, globalRot, localPos, localRot);

	dBodyID body;
	if(entity->getStill()) body = 0;
	else {
		body = dBodyCreate(dynWorld);
		bodies[entity] = body;
	
		dBodySetPosition(body, newGlobalPos[0], newGlobalPos[1], newGlobalPos[2]);
		dBodySetRotation(body, newGlobalRot);
	}
		    
	dMass mass;
	dMassSetZero(&mass);
	vector<Shape*>& shapes = entity->getShapes();
	vector<dGeomID> geoms;

	for(vector<Shape*>::iterator shape = shapes.begin(); shape != shapes.end(); shape++) {
		dGeomID geom = initShape(entity, body, Physics::colSpace, &mass, *shape, newGlobalPos, newGlobalRot);
		geoms.push_back(geom);
	}
	
	vector<Sensor*>& sensors = entity->getSensors();
	for(vector<Sensor*>::iterator sensor = sensors.begin(); sensor != sensors.end(); sensor++) {
		GenericPhysicalSensor *physicalSensor = dynamic_cast<GenericPhysicalSensor*>(*sensor);
		physicalSensor->init(body);
	}
	
//	cout << "mass center: " << mass.c[0] << ", " << mass.c[1] << ", " << mass.c[2] << endl;
	if(mass.mass!=0 && !entity->getStill()) {
		if(mass.c[0] != 0 || mass.c[1] != 0 || mass.c[2] != 0) {
			// special case - we must move center of mass to 0,0,0 to let it be handled correctly
			// to do it, we are shifting geoms and body
			// based on: http://opende.sourceforge.net/wiki/index.php/HOWTO_composite_objects
			for(vector<dGeomID>::iterator geom = geoms.begin(); geom != geoms.end(); geom++) {
				const dReal* position = dGeomGetPosition(*geom);
				dGeomSetPosition(*geom, position[0]-mass.c[0], position[1]-mass.c[1], position[2]-mass.c[2]);
			}
	
			// compute position shift for special case with composite bodies
			shifts[body] = Vector3D(-mass.c[0], -mass.c[1], -mass.c[2]);
	
		    dVector3 res;
		    dMULTIPLY0_331 (res,localRot,mass.c);	
//			cout << "moving position by " << res[0] << ", " << res[1] << ", " << res[2] << endl;
			dBodySetPosition(body, newGlobalPos[0]+res[0], newGlobalPos[1]+res[1], newGlobalPos[2]+res[2]);
	
			dMassTranslate(&mass, -mass.c[0], -mass.c[1], -mass.c[2]);
		}
		dBodySetMass(body, &mass);
	}
	
	vector<Joint*>& joints = entity->getJoints();
	for(vector<Joint*>::iterator joint = joints.begin(); joint != joints.end(); joint++) {
		dBodyID subBody = initEntity((*joint)->getEntity(), newGlobalPos, newGlobalRot);
		initJoint(*joint, body, subBody, newGlobalPos, newGlobalRot);
	}
	return body;
}

void Physics::init(World* root) {
	Physics::root = root;
	dInitODE();
	dynWorld = dWorldCreate();
	dWorldSetGravity(dynWorld, 0, -9.81, 0);
	dWorldSetERP(dynWorld, 0.25);
	dWorldSetCFM(dynWorld, 1e-4);
	colSpace = dHashSpaceCreate(0);
	vector<Entity*>& entities = root->getEntities();
	dVector3 globalPos;
	dSetZero(globalPos, 3);
	dMatrix3 globalRot;
	dRSetIdentity(globalRot);
	
	for(vector<Entity*>::iterator entity = entities.begin(); entity != entities.end(); entity++) {
		initEntity(*entity, globalPos, globalRot);
	}
	contactGroup = dJointGroupCreate(0);
}

void Physics::nearCallback (void *data, dGeomID o1, dGeomID o2) {
	dContactGeom contactGeoms[10];
	int max_contacts = sizeof(contactGeoms) / sizeof(contactGeoms[0]);
	if (dGeomIsSpace (o1) || dGeomIsSpace (o2)) {
		// colliding a space with something
		dSpaceCollide2 (o1,o2,data,&nearCallback);
		// collide all geoms internal to the space(s)
		if (dGeomIsSpace (o1)) dSpaceCollide ((dxSpace*)o1,data,&nearCallback);
		if (dGeomIsSpace (o2)) dSpaceCollide ((dxSpace*)o2,data,&nearCallback);
	} else {
		counter++;
		
		dBodyID body1 = dGeomGetBody(o1);
		dBodyID body2 = dGeomGetBody(o2);

		if (body1 == body2 || !(
			jointBodies.find(std::make_pair(body1, body2)) == jointBodies.end() && 
			jointBodies.find(std::make_pair(body2, body1)) == jointBodies.end()
		)) return;

		// colliding two non-space geoms, so generate contact
		// points between o1 and o2
		int num_contact = dCollide(o1,o2,max_contacts,contactGeoms,sizeof(contactGeoms[0]));

//		cout << "callback! " << body1 << ", " << body2 << ", " << num_contact << endl;
		
		if(num_contact > 0) {
//			std::cout << "collision! " << body1 << ", " << body2 << std::endl;
			Shape *s1;
			if(dGeomGetClass(o1) == dGeomTransformClass) s1 = (Shape*) dGeomGetData(dGeomTransformGetGeom(o1));
			else s1 = (Shape*) dGeomGetData(o1);
			
			Shape *s2;
			if(dGeomGetClass(o2) == dGeomTransformClass) s2 = (Shape*) dGeomGetData(dGeomTransformGetGeom(o2));
			else s2 = (Shape*) dGeomGetData(o2);
			
			// add these contact points to the simulation
			for(int i=0; i<num_contact; i++) {
		      	dContact contactInfo;
				contactInfo.geom = contactGeoms[i];
				contactInfo.surface.mode = 0;

				// Compute the CFM and ERP by assuming the two bodies form a
				// spring-damper system.
				double h, kp, kd;
				h = root->getStep();
				kp = 1 / (1 / s1->getMaterial()->getResilience() + 1 / s2->getMaterial()->getResilience());
				kd = s1->getMaterial()->getDamping() + s2->getMaterial()->getDamping();
				contactInfo.surface.mode |= dContactSoftERP | dContactSoftCFM;
				contactInfo.surface.soft_erp = h * kp / (h * kp + kd);
				contactInfo.surface.soft_cfm = 1 / (h * kp + kd);

				// Compute friction effects; this is standard Coulomb friction
				contactInfo.surface.mode |= dContactApprox1;
				contactInfo.surface.mu = std::min(s1->getMaterial()->getFriction(), s2->getMaterial()->getFriction());
			
				// Construct a contact joint between the two bodies
				dJointID joint = dJointCreateContact(dynWorld, contactGroup, &contactInfo);
				dJointFeedback *feedback = new dJointFeedback;
				dJointSetFeedback(joint, feedback);
				feedbacks[feedback] = std::make_pair(s1, s2);
				dJointAttach(joint, body1, body2);
			}
		}
	}
}
  
void Physics::step() {
	feedbacks.clear();
	dSpaceCollide(colSpace,0,&nearCallback);
//	cout << "counter " << counter << endl;
	counter = 0;
	dWorldQuickStep(dynWorld, root->getStep());
	forces.clear();
	for(map<dJointFeedback*, pair<Shape*, Shape*> >::iterator feedback = feedbacks.begin(); feedback != feedbacks.end(); feedback++) {
    	dJointFeedback *jointFeedback = feedback->first;
		if(forces.find(feedback->second.first) == forces.end()) {
			forces[feedback->second.first] = Vector3D(jointFeedback->f1[0],jointFeedback->f1[1],jointFeedback->f1[2]);
    	} else {
			forces[feedback->second.first].x += jointFeedback->f1[0];
			forces[feedback->second.first].y += jointFeedback->f1[1];
			forces[feedback->second.first].z += jointFeedback->f1[2];
		}
		if(forces.find(feedback->second.second) == forces.end()) {
			forces[feedback->second.second] = Vector3D(jointFeedback->f2[0],jointFeedback->f2[1],jointFeedback->f2[2]);
    	} else {
			forces[feedback->second.second].x += jointFeedback->f2[0];
			forces[feedback->second.second].y += jointFeedback->f2[1];
			forces[feedback->second.second].z += jointFeedback->f2[2];
		}
		delete jointFeedback;
	}
	feedbacks.clear();
	
	for(map<Joint*, dJointID>::iterator it = joints.begin(); it != joints.end(); it++) {
		GenericPhysicalJoint* joint = dynamic_cast<GenericPhysicalJoint*>(it->first);
		dJointID jointID = it->second;

		joint->update(jointID);
	}    
 
 	dVector3 globalPos;
	dSetZero(globalPos, 3);
	dMatrix3 globalRot;
	dRSetIdentity(globalRot);
 	vector<Entity*>& entities = root->getEntities();
	for(vector<Entity*>::iterator entity = entities.begin(); entity != entities.end(); entity++) {
		updateEntity(*entity, globalPos, globalRot);
	}
    
    dJointGroupEmpty(contactGroup);
}

void Physics::cleanup() {
	dWorldDestroy(dynWorld);
}

void Physics::convertVector(Vector3D& src, dVector3 dest) {
	dest[0] = src.x;
	dest[1] = src.y;
	dest[2] = src.z;
}

void Physics::convertQuaternion(Quaternion& src, dQuaternion dest) {
	dest[0] = src.w;
	dest[1] = src.x;
	dest[2] = src.y;
	dest[3] = src.z;
}

void Physics::updateEntity(Entity* entity, dVector3 parentGlobalPos, dMatrix3 parentGlobalRot) {
	dVector3 globalPos;
	dMatrix3 globalRot;
	
	if(bodies.find(entity) != bodies.end()) {
		// Entity with body, we have to compute local body pose
		// and pass its global position and rotation to Entities in subtree rooted in current Entity
		dBodyID body = bodies[entity];
		dReal *position = (dReal*) dBodyGetPosition(body);
		dReal *q = (dReal*) dBodyGetQuaternion(body);
	 	if(shifts.find(body)!=shifts.end()) {
			// special case with shifted position for composite bodies
			dVector3 shift;
			convertVector(shifts[body], shift);
		    dVector3 res;
			dMatrix3 globalRot;
			dQtoR(q, globalRot);
		    dMULTIPLY0_331 (res,globalRot,shift);
		    dOP(globalPos,+,position,res);	
		} else {
			globalPos[0] = position[0];
			globalPos[1] = position[1];
			globalPos[2] = position[2];
		}
		dQtoR(q, globalRot);
		
		dMatrix3 parentGlobalRotInverted;

		dQuaternion parentGlobalRotQuaternion;
		dRtoQ(parentGlobalRot, parentGlobalRotQuaternion);
		dQuaternion parentGlobalRotInvertedQuaternion;
		parentGlobalRotInvertedQuaternion[0] = -parentGlobalRotQuaternion[0];
		parentGlobalRotInvertedQuaternion[1] = parentGlobalRotQuaternion[1];
		parentGlobalRotInvertedQuaternion[2] = parentGlobalRotQuaternion[2];
		parentGlobalRotInvertedQuaternion[3] = parentGlobalRotQuaternion[3];
		dQtoR(parentGlobalRotInvertedQuaternion, parentGlobalRotInverted);
		
		dVector3 localPosRotated;
		dOP(localPosRotated,-,globalPos,parentGlobalPos);
		dVector3 localPos;
		dMULTIPLY0_331(localPos,parentGlobalRotInverted,localPosRotated);
		dMatrix3 localRot;
		dMULTIPLY0_333(localRot,parentGlobalRotInverted,globalRot);
		
		dQuaternion localRotQ;
		dRtoQ(localRot, localRotQ);

		Quaternion quat;
		quat.w = localRotQ[0];
		quat.x = localRotQ[1];
		quat.y = localRotQ[2];
		quat.z = localRotQ[3];
    	entity->setRotation(quat);
    	entity->setPosition(localPos[0], localPos[1], localPos[2]);
	} else {
		// static Entity without body, we don't have to set body pose in that case, only
		// update global position and rotation of Entities in subtree rooted in current Entity
		
		Vector3D _localPos = entity->getPosition();
		Quaternion& _localRot = entity->getRotation();
		dVector3 localPos;
		convertVector(_localPos, localPos);
		dQuaternion localRotQ;
		convertQuaternion(_localRot, localRotQ);
		dMatrix3 localRot;
		dQtoR(localRotQ, localRot);
		
		getGlobalPose(globalPos, globalRot, parentGlobalPos, parentGlobalRot, localPos, localRot);
	}
	
	vector<Sensor*>& sensors = entity->getSensors();
	for(vector<Sensor*>::iterator sensor = sensors.begin(); sensor != sensors.end(); sensor++) {
		GenericPhysicalSensor *physicalSensor = dynamic_cast<GenericPhysicalSensor*>(*sensor);
		physicalSensor->update();
	}	
	
	vector<Joint*>& joints = entity->getJoints();
	for(vector<Joint*>::iterator joint = joints.begin(); joint != joints.end(); joint++) {
		updateEntity((*joint)->getEntity(), globalPos, globalRot);
	}
}
